import argparse
from pathlib import Path
import torch
from tqdm import tqdm
import os 
os.environ['CUDA_VISIBLE_DEVICES'] = "1" 

from pcdet.config import cfg, cfg_from_yaml_file
from pcdet.datasets import build_dataloader
from pcdet.models import build_network, load_data_to_gpu
from pcdet.utils import common_utils

def parse_config():
    parser = argparse.ArgumentParser(description='arg parser')
    parser.add_argument('--cfg_file',
                        type=str,
                        default='cfgs/kitti_models/second.yaml',
                        help='specify the config for demo')
    parser.add_argument('--ckpt',
                        type=str,
                        default=None,
                        help='specify the pretrained model')
    parser.add_argument('--batch_size',
                        type=int,
                        default=2,
                        required=False,
                        help='batch size for inference')
    parser.add_argument('--workers',
                        type=int,
                        default=0,
                        help='number of workers for dataloader')
    parser.add_argument('--extra_tag', type=str, default='default', help='extra tag for this experiment')
    args = parser.parse_args()

    cfg_from_yaml_file(args.cfg_file, cfg)
    cfg.TAG = Path(args.cfg_file).stem
    cfg.EXP_GROUP_PATH = '/'.join(
        args.cfg_file.split('/')[1:-1])  # remove 'cfgs' and 'xxxx.yaml'
    cfg.EPOCH_ID = Path(args.ckpt).name[17:-4]

    return args, cfg


def main():
    args, cfg = parse_config()
    logger = common_utils.create_logger()
    logger.info(
        '-----------------Inference of OpenPCDet-------------------------')
    test_set, test_loader, sampler = build_dataloader(
        dataset_cfg=cfg.DATA_CONFIG,
        class_names=cfg.CLASS_NAMES,
        batch_size=args.batch_size,
        dist=False,
        workers=args.workers,
        logger=logger,
        training=False)
    logger.info(f'Total number of samples: \t{len(test_set)}')

    model = build_network(model_cfg=cfg.MODEL,
                          num_class=len(cfg.CLASS_NAMES),
                          dataset=test_set)
    model.load_params_from_file(filename=args.ckpt, logger=logger, to_cpu=True)
    model.cuda()
    model.eval()
    output_dir = cfg.ROOT_DIR / 'output' / cfg.EXP_GROUP_PATH / cfg.TAG / args.extra_tag
    output_dir.mkdir(parents=True, exist_ok=True)
    inf_output_dir = output_dir / 'inference' / cfg.EPOCH_ID
    inf_output_dir.mkdir(parents=True, exist_ok=True) 
    for _, batch_dict in enumerate(tqdm(test_loader)):
        load_data_to_gpu(batch_dict)
        with torch.no_grad():
            pred_dicts, _ = model(batch_dict)
            annos = test_set.generate_prediction_dicts(batch_dict,
                                                       pred_dicts,
                                                       cfg.CLASS_NAMES,
                                                       output_path=inf_output_dir)

if __name__ == '__main__':
    main()